{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "import argparse\n",
    "import copy\n",
    "import json\n",
    "import os\n",
    "import pickle\n",
    "import random\n",
    "import sys\n",
    "import re\n",
    "import numpy as np\n",
    "import openai\n",
    "from typing import List, Dict, Any\n",
    "import ast\n",
    "from p3_utils import *\n",
    "from tqdm import tqdm\n",
    "\n",
    "#GPT_MODEL = \"gpt-4-turbo-preview\"  # [gpt-4-turbo-preview]\n",
    "GPT_MODEL = \"gpt-3.5-turbo\"  # [gpt-3.5-turbo-preview]\n",
    "\n",
    "setOpenAi(keyid =4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cot_solve(question):\n",
    "    cot_prompt = \"\"\"\n",
    "    You will be given a programming puzzle.\n",
    "    Your task is to find an input for the following function that will make the function return True. \n",
    "    Let's think step by step\"\"\"\n",
    "    \n",
    "    input_q = [\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": cot_prompt\n",
    "        },\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": question\n",
    "        }\n",
    "    ]\n",
    "    \n",
    "    procedure = askChatGPT(input_q, model = GPT_MODEL, temperature=1)\n",
    "    ans_procedure = {\"role\": \"assistant\", \"content\": procedure}\n",
    "    input_q.append(ans_procedure)\n",
    "    #final_q = {\"role\": \"user\", \"content\": \"So, what is the input? You can either give the input or write a solution function in Python. Make sure the answer you return in the format 'def function_name(): return value'. Make sure NO explanation!\"}\n",
    "    final_q = {\"role\": \"user\", \"content\": \"\"\"So, what is the input? Please only give the input in the format of a string and just give the answer without any additional explanation or clarification, no prefix or suffix.\n",
    "\n",
    "    For example if the input should be x = 5, then you should only give the answer as 5 and not x = 5.\n",
    "    For example, if the the input is list = [1,2,3], then you should only give the answer as [1,2,3] and not list = [1,2,3].\n",
    "                \"\"\"}\n",
    "    input_q.append(final_q)\n",
    "    final_answer = askChatGPT(input_q, model = GPT_MODEL, temperature=1)\n",
    "    return final_answer\n",
    "\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "puzzles = []\n",
    "\n",
    "\n",
    "with open(\"/Users/natehu/Desktop/Tsinghua Research/TaDe/dataset_gen_using_cot/p3_100_test_questions.json\", \"r\") as f:\n",
    "    puzzles = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "id = 33\n",
    "\n",
    "question = puzzles[id][\"sat\"]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "solution = cot_solve(question)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'-5.770237138115334'"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "solution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "converted_result = convert_to_type(puzzles[id]['ans_type'], solution)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "exec(question)\n",
    "result = sat(converted_result)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_price(number_of_token, type, model):\n",
    "    if model == \"gpt-3.5-turbo\":\n",
    "        if type == \"input\":\n",
    "            per_token = 0.5 / 1000000\n",
    "            return number_of_token * per_token\n",
    "        elif type == \"output\":\n",
    "            per_token = 1.5 / 1000000\n",
    "            return number_of_token * per_token\n",
    "    elif model == \"gpt-4-turbo-preview\":\n",
    "        if type == \"input\":\n",
    "            per_token = 10 / 1000000\n",
    "            return number_of_token * per_token\n",
    "        elif type == \"output\":\n",
    "            per_token = 30 / 1000000\n",
    "            return number_of_token * per_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.20243999999999998"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "calculate_price(22606, \"output\", \"gpt-3.5-turbo\") + calculate_price(337062, \"input\", \"gpt-3.5-turbo\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
